import os
import xml.etree.ElementTree as ET
from getclass import doit

# 指定数据集所在目录
data_dir = 'D:/work/annotations/xmls'
output_dir = 'D:/work/test'
# 定义标签名称与对应的索引
LABELS = doit(data_dir)

# 遍历数据集目录中的每一个XML文件并转换为TXT格式
for filename in os.listdir(data_dir):
    if not filename.endswith('.xml'):
        continue
    lens = len(filename)
    name = ""
    for i in range(1, lens + 1):
        if filename[-i] == '_':
            name = filename[:lens - i]
            break
    xml_path = os.path.join(data_dir, filename)
    tree = ET.parse(xml_path)
    root = tree.getroot()
    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)

    with open(os.path.join(output_dir, filename.replace('.xml', '.txt')), 'w') as f:
        for obj in root.findall('object'):
            label = name
            if label not in LABELS.keys():
                continue
            label_idx = LABELS[label]
            bbox = obj.find('bndbox')
            xmin = int(bbox.find('xmin').text)
            ymin = int(bbox.find('ymin').text)
            xmax = int(bbox.find('xmax').text)
            ymax = int(bbox.find('ymax').text)
            x_center = (xmin + xmax) / (2.0 * width)
            y_center = (ymin + ymax) / (2.0 * height)
            w = (xmax - xmin) / float(width)
            h = (ymax - ymin) / float(height)
            f.write(f"{label_idx} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}\n")
